from typing import Iterable, Iterator, TypeVar, TypedDict

from zkl_aiutils_datasets import Dataset, DatasetIterator, IndexedIterator, skip, wrap_dataset_iterator

AnyItem = TypeVar('AnyItem')


class ChannelsBatchIteratorState(TypedDict):
    channels_seq_i: tuple[int | None, ...]
    channels_item_i: tuple[int | None, ...]


class ChannelsBatchedIterator(DatasetIterator[tuple[AnyItem, ...]]):
    NoPadding = object()

    def __init__(self,
        sequences: Iterable[Iterable[AnyItem]],
        channels_n: int,
        pad_item: AnyItem = NoPadding,
        state: ChannelsBatchIteratorState | None = None,
    ):
        self._sequences_iterator = IndexedIterator(sequences)
        self._channels_n = channels_n
        self._pad_item = pad_item

        self._channels_seq_i: list[int | None] = []
        self._channels_item_i: list[int | None] = []
        self._channels_items_iterator: list[Iterator[AnyItem] | None] = []

        if state is not None:
            self._channels_seq_i = list(state['channels_seq_i'])
            self._channels_item_i = list(state['channels_item_i'])
            self._channels_items_iterator = [None] * len(self._channels_seq_i)

            if any((channel_i is None) for channel_i in self._channels_seq_i):
                if self._pad_item is self.NoPadding:
                    raise ValueError('Cannot resume ChannelsBatchedIterator without pad_item!')

            indexed_channels_seq_i = enumerate(self._channels_seq_i)
            filtered_channels_seq_i = filter(lambda args: args[1] is not None, indexed_channels_seq_i)
            sorted_channels_seq_i = sorted(filtered_channels_seq_i, key=lambda args: args[1])

            last_seq_i = 0
            for channel_i, seq_i in sorted_channels_seq_i:
                skip(self._sequences_iterator, seq_i - last_seq_i)
                items_iterator, _ = next(self._sequences_iterator)
                items_iterator = wrap_dataset_iterator(items_iterator)
                skip(items_iterator, self._channels_item_i[channel_i])
                self._channels_items_iterator[channel_i] = items_iterator
                last_seq_i = seq_i + 1

    def __next__(self) -> tuple[AnyItem, ...]:
        while len(self._channels_seq_i) < self._channels_n:
            try:
                items_iterator, seq_i = next(self._sequences_iterator)
                items_iterator = wrap_dataset_iterator(items_iterator)
                self._channels_seq_i.append(seq_i)
                self._channels_item_i.append(0)
                self._channels_items_iterator.append(items_iterator)
            except StopIteration:
                if self._pad_item is self.NoPadding:
                    raise
                self._channels_seq_i.append(None)
                self._channels_item_i.append(None)
                self._channels_items_iterator.append(None)

        if all((iterator is None) for iterator in self._channels_items_iterator):
            raise StopIteration

        channels_item = []
        for channel_i in range(self._channels_n):
            while True:
                try:
                    iterator = self._channels_items_iterator[channel_i]
                    if iterator is not None:
                        channels_item.append(next(iterator))
                        self._channels_item_i[channel_i] += 1
                    else:
                        channels_item.append(self._pad_item)
                    break
                except StopIteration:
                    if len(self._channels_seq_i) > self._channels_n:
                        self._channels_seq_i[channel_i] = self._channels_seq_i.pop()
                        self._channels_item_i[channel_i] = self._channels_item_i.pop()
                        self._channels_items_iterator[channel_i] = self._channels_items_iterator.pop()
                        if all((iterator is None) for iterator in self._channels_items_iterator):
                            raise StopIteration
                    else:
                        try:
                            items_iterator, seq_i = next(self._sequences_iterator)
                            items_iterator = wrap_dataset_iterator(items_iterator)
                            self._channels_seq_i[channel_i] = seq_i
                            self._channels_item_i[channel_i] = 0
                            self._channels_items_iterator[channel_i] = items_iterator
                        except StopIteration:
                            if self._pad_item is self.NoPadding:
                                raise
                            self._channels_seq_i[channel_i] = None
                            self._channels_item_i[channel_i] = None
                            self._channels_items_iterator[channel_i] = None
                            if all((iterator is None) for iterator in self._channels_items_iterator):
                                raise StopIteration
                    continue
        return tuple(channels_item)

    def __pause__(self) -> ChannelsBatchIteratorState:
        return ChannelsBatchIteratorState(
            channels_seq_i=tuple(self._channels_seq_i),
            channels_item_i=tuple(self._channels_item_i))


class ChannelsBatchedDataset(Dataset[tuple[AnyItem, ...]]):
    NoPadding = ChannelsBatchedIterator.NoPadding

    def __init__(self,
        sequences: Iterable[Iterable[AnyItem]],
        channels_n: int,
        pad_item: AnyItem = NoPadding,
    ):
        self._sequences = sequences
        self._channels_n = channels_n
        self._pad_item = pad_item

    @property
    def channels_n(self) -> int:
        return self._channels_n

    @property
    def sequences(self) -> Iterable[Iterable[AnyItem]]:
        return self._sequences

    def __iter__(self) -> ChannelsBatchedIterator[AnyItem]:
        return ChannelsBatchedIterator(self._sequences, self._channels_n, self._pad_item)

    def __resume__(self, state: ChannelsBatchIteratorState | None = None) -> ChannelsBatchedIterator[AnyItem]:
        return ChannelsBatchedIterator(self._sequences, self._channels_n, self._pad_item, state)
